# apps/orchestrator/task_review_manager.py
from typing import Dict, Any, List, Optional
from datetime import datetime
import json
from loguru import logger
from collections import Counter
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 导入核心依赖
from utils.logger import Logger
from utils.exceptions import RedisError, TaskReviewError, ModelError
from states.task_state import TaskStateManager
from middleware.redis_adapter import get_redis_adapter
from middleware.model_interface import get_model_client, BaseModelClient
from config.model_config import ModelBackend, get_model_config

class TaskReviewManager:
    """任务回顾管理器（负责智能体完成子任务后的结果总结和冗余信息过滤）"""
    
    def __init__(self, redis_client=None):
        self.redis_client = redis_client or get_redis_adapter()
        self.logger = Logger.get_logger(agent_name="task_review_manager")
        self.model_config = get_model_config()
        self.model_client: Optional[BaseModelClient] = None
        self.model_backend = None
    
    async def _init_model_client(self, model_backend: Optional[str] = None):
        """初始化模型客户端，支持多种后端"""
        if self.model_client is not None and self.model_backend == model_backend:
            return
        
        self.model_backend = model_backend
        
        # 尝试初始化模型客户端
        try:
            if model_backend:
                # 将字符串转换为ModelBackend枚举
                backend_enum = ModelBackend(model_backend)
                self.model_client = get_model_client(backend=backend_enum)
                self.logger.info(f"初始化{backend_enum.value}模型客户端成功")
            else:
                # 尝试按优先级初始化可用后端
                for backend in [ModelBackend.VLLM, ModelBackend.OLLAMA]:
                    try:
                        self.model_client = get_model_client(backend=backend)
                        self.model_backend = backend.value
                        self.logger.info(f"初始化{backend.value}模型客户端成功")
                        return
                    except Exception as e:
                        self.logger.warning(f"初始化{backend.value}模型客户端失败: {str(e)}")
                
                self.logger.warning("所有模型后端初始化失败，将使用规则降级模式")
                self.model_client = get_model_client(backend=ModelBackend.FALLBACK)
                self.model_backend = ModelBackend.FALLBACK.value
                
        except Exception as e:
            self.logger.error(f"初始化模型客户端异常: {str(e)}")
            self.model_client = get_model_client(backend=ModelBackend.FALLBACK)
            self.model_backend = ModelBackend.FALLBACK.value
    
    async def review_task_results(self, task_id: str, model_backend: Optional[str] = None) -> Dict[str, Any]:
        """
        回顾任务执行结果
        
        Args:
            task_id: 任务ID
            model_backend: 可选的模型后端
            
        Returns:
            包含总结结果和过滤后信息的字典
        """
        try:
            self.logger.info(f"开始回顾任务结果: task_id={task_id}, backend={model_backend}")
            
            # 获取任务状态
            task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
            task_state = task_manager.get_state()
            
            # 获取所有子任务结果
            sub_tasks_results = await self._get_sub_tasks_results(task_id)
            
            if not sub_tasks_results:
                self.logger.warning(f"任务 {task_id} 没有子任务结果")
                return {
                    "task_id": task_id,
                    "summary": "暂无子任务结果",
                    "filtered_results": [],
                    "redundancy_rate": 0.0,
                    "status": "empty"
                }
            
            # 过滤冗余信息
            filtered_results = self._filter_redundant_information(sub_tasks_results)
            
            # 计算冗余率
            redundancy_rate = 1.0 - (len(filtered_results) / len(sub_tasks_results)) if sub_tasks_results else 0.0
            
            # 生成总结
            summary = await self._generate_summary(filtered_results, model_backend)
            
            # 保存回顾结果
            await self._save_review_result(task_id, summary, filtered_results, redundancy_rate)
            
            self.logger.info(f"任务回顾完成: task_id={task_id}")
            
            return {
                "task_id": task_id,
                "summary": summary,
                "filtered_results": filtered_results,
                "redundancy_rate": round(redundancy_rate, 2),
                "status": "completed"
            }
        
        except Exception as e:
            self.logger.error(f"回顾任务结果失败: task_id={task_id}, error={str(e)}")
            raise TaskReviewError(
                message="Task review failed", 
                context={"task_id": task_id, "error": str(e)}
            ) from e
    
    async def _get_sub_tasks_results(self, task_id: str) -> List[Dict[str, Any]]:
        """获取任务的所有子任务结果"""
        try:
            # 从Redis获取子任务结果
            sub_tasks_key = f"task:{task_id}:sub_tasks"
            sub_task_ids = self.redis_client.smembers(sub_tasks_key)
            
            sub_tasks_results = []
            for sub_task_id in sub_task_ids:
                sub_task_manager = TaskStateManager(task_id=sub_task_id.decode() if isinstance(sub_task_id, bytes) else sub_task_id)
                sub_task_state = sub_task_manager.get_state()
                
                if sub_task_state["state"] == "success" and sub_task_state["output_result"]:
                    sub_tasks_results.append({
                        "sub_task_id": sub_task_state["task_id"],
                        "agent_type": sub_task_state["agent_type"],
                        "agent_instance_id": sub_task_state["agent_instance_id"],
                        "result": sub_task_state["output_result"],
                        "created_at": sub_task_state["created_at"],
                        "updated_at": sub_task_state["updated_at"]
                    })
            
            # 按更新时间排序
            sub_tasks_results.sort(key=lambda x: x["updated_at"])
            
            return sub_tasks_results
        except Exception as e:
            self.logger.error(f"获取子任务结果失败: task_id={task_id}, error={str(e)}")
            return []
    
    def _filter_redundant_information(self, sub_tasks_results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """过滤冗余信息"""
        if len(sub_tasks_results) <= 1:
            return sub_tasks_results
        
        # 提取文本内容
        texts = []
        for result in sub_tasks_results:
            # 尝试从不同格式的结果中提取文本
            if isinstance(result["result"], dict):
                text = self._extract_text_from_result(result["result"])
            else:
                text = str(result["result"])
            texts.append(text)
        
        # 使用TF-IDF和余弦相似度检测冗余
        filtered_results = []
        seen_indices = set()
        
        # 如果文本数量大于1，进行相似度计算
        if len(texts) > 1:
            # 创建TF-IDF向量
            vectorizer = TfidfVectorizer(stop_words='english', max_features=1000)
            try:
                tfidf_matrix = vectorizer.fit_transform(texts)
                
                # 计算相似度矩阵
                similarity_matrix = cosine_similarity(tfidf_matrix)
                
                # 过滤冗余（相似度大于0.8的视为冗余）
                for i in range(len(sub_tasks_results)):
                    if i in seen_indices:
                        continue
                    
                    filtered_results.append(sub_tasks_results[i])
                    
                    # 标记相似的结果为已处理
                    for j in range(i + 1, len(sub_tasks_results)):
                        if similarity_matrix[i, j] > 0.8 and j not in seen_indices:
                            seen_indices.add(j)
                            self.logger.debug(f"过滤冗余结果: {sub_tasks_results[j]['sub_task_id']}")
            except Exception as e:
                self.logger.warning(f"TF-IDF相似度计算失败，使用简单文本相似度: {str(e)}")
                # 回退到简单的文本相似度计算
                filtered_results = self._simple_text_similarity_filter(sub_tasks_results, texts)
        else:
            filtered_results = sub_tasks_results
        
        return filtered_results
    
    def _simple_text_similarity_filter(self, sub_tasks_results: List[Dict[str, Any]], texts: List[str]) -> List[Dict[str, Any]]:
        """简单的文本相似度过滤（备用方法）"""
        filtered_results = []
        seen_indices = set()
        
        for i in range(len(texts)):
            if i in seen_indices:
                continue
                
            filtered_results.append(sub_tasks_results[i])
            
            # 使用词频比较
            words1 = self._get_words(texts[i])
            word_count1 = Counter(words1)
            
            for j in range(i + 1, len(texts)):
                if j in seen_indices:
                    continue
                    
                words2 = self._get_words(texts[j])
                word_count2 = Counter(words2)
                
                # 计算Jaccard相似度
                intersection = sum((word_count1 & word_count2).values())
                union = sum((word_count1 | word_count2).values())
                similarity = intersection / union if union > 0 else 0
                
                if similarity > 0.7:  # 相似度阈值
                    seen_indices.add(j)
                    self.logger.debug(f"简单过滤冗余结果: {sub_tasks_results[j]['sub_task_id']}")
        
        return filtered_results
    
    def _get_words(self, text: str) -> List[str]:
        """从文本中提取单词"""
        # 简单的单词提取
        import re
        words = re.findall(r'\b\w+\b', text.lower())
        return words
    
    def _extract_text_from_result(self, result: Dict[str, Any]) -> str:
        """从结果字典中提取文本"""
        # 尝试从常见的字段中提取文本
        text_fields = ['answer', 'result', 'content', 'text', 'summary', 'description']
        
        for field in text_fields:
            if field in result:
                if isinstance(result[field], str):
                    return result[field]
                elif isinstance(result[field], dict):
                    return self._extract_text_from_result(result[field])
                elif isinstance(result[field], list):
                    # 提取列表中的文本
                    texts = []
                    for item in result[field]:
                        if isinstance(item, dict):
                            texts.append(self._extract_text_from_result(item))
                        elif isinstance(item, str):
                            texts.append(item)
                    return ' '.join(texts)
        
        # 如果没有找到合适的文本字段，返回整个结果的字符串表示
        return str(result)
    
    async def _generate_summary(self, filtered_results: List[Dict[str, Any]], model_backend: Optional[str] = None) -> str:
        """生成任务结果总结"""
        if not filtered_results:
            return "没有可用的任务结果"
        
        # 提取关键信息
        key_infos = []
        for result in filtered_results:
            if isinstance(result["result"], dict):
                text = self._extract_text_from_result(result["result"])
            else:
                text = str(result["result"])
            
            key_infos.append({
                "agent_type": result["agent_type"],
                "text": text
            })
        
        # 尝试使用模型生成总结
        summary = await self._generate_summary_with_model(key_infos, model_backend)
        
        # 如果模型生成失败，使用简单的总结方法
        if not summary:
            summary = self._simple_summary(key_infos)
        
        return summary
    
    async def _generate_summary_with_model(self, key_infos: List[Dict[str, str]], model_backend: Optional[str] = None) -> Optional[str]:
        """使用模型生成总结"""
        try:
            # 初始化模型客户端
            await self._init_model_client(model_backend)
            
            if not self.model_client or self.model_backend == ModelBackend.FALLBACK.value:
                self.logger.warning("模型客户端不可用，将使用规则降级模式")
                return None
            
            # 构建提示
            system_prompt = "你是一个专业的任务总结助手。请根据多个子任务的执行结果，生成一个全面、简洁的总结。"
            user_prompt = self._build_summary_prompt(key_infos)
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
            
            # 使用统一接口生成总结
            summary = await self.model_client.chat_completion(
                messages=messages,
                max_tokens=500,
                temperature=0.3
            )
            
            return summary
            
        except Exception as e:
            self.logger.warning(f"使用模型生成总结失败: {str(e)}")
            return None
    
    def _build_summary_prompt(self, key_infos: List[Dict[str, str]]) -> str:
        """构建总结提示"""
        prompt = "请基于以下智能体执行结果，生成一个简洁、连贯的总结。总结应该包含所有关键信息，并避免冗余。\n\n"
        
        for i, info in enumerate(key_infos, 1):
            prompt += f"智能体类型: {info['agent_type']}\n"
            prompt += f"结果 {i}: {info['text']}\n\n"
        
        prompt += "请生成总结:"
        return prompt
    
    def _simple_summary(self, key_infos: List[Dict[str, str]]) -> str:
        """简单的总结方法（备用）"""
        if not key_infos:
            return "没有可用的结果"
        
        summary = "任务执行结果总结:\n\n"
        
        # 统计不同智能体类型
        agent_types = Counter([info["agent_type"] for info in key_infos])
        
        summary += f"- 共有 {len(key_infos)} 个子任务执行结果\n"
        summary += f"- 涉及智能体类型: {', '.join([f'{k}({v})' for k, v in agent_types.items()])}\n\n"
        
        # 添加每个智能体的主要结果
        for i, info in enumerate(key_infos[:5], 1):  # 最多显示5个结果
            summary += f"{i}. [{info['agent_type']}] {self._truncate_text(info['text'], 100)}\n"
        
        if len(key_infos) > 5:
            summary += f"... 等 {len(key_infos) - 5} 个结果\n"
        
        return summary
    
    def _truncate_text(self, text: str, max_length: int) -> str:
        """截断文本"""
        if len(text) <= max_length:
            return text
        return text[:max_length-3] + "..."
    
    async def _save_review_result(self, task_id: str, summary: str, filtered_results: List[Dict[str, Any]], redundancy_rate: float) -> bool:
        """保存回顾结果"""
        try:
            review_data = {
                "task_id": task_id,
                "summary": summary,
                "filtered_results": json.dumps([r["sub_task_id"] for r in filtered_results], ensure_ascii=False),
                "redundancy_rate": str(redundancy_rate),
                "created_at": str(int(datetime.now().timestamp()))
            }
            
            # 保存到Redis
            review_key = f"task:{task_id}:review"
            self.redis_client.hsetall(review_key, review_data)
            self.redis_client.expire(review_key, 86400)  # 保存24小时
            
            # 更新任务状态，添加回顾信息
            task_manager = TaskStateManager(task_id=task_id, redis_client=self.redis_client)
            task_state = task_manager.get_state()
            
            # 合并原有输出结果和回顾信息
            output_result = task_state.get("output_result", {})
            if isinstance(output_result, str):
                try:
                    output_result = json.loads(output_result)
                except:
                    output_result = {}
            
            output_result["review_summary"] = summary
            output_result["redundancy_rate"] = redundancy_rate
            output_result["filtered_sub_task_ids"] = [r["sub_task_id"] for r in filtered_results]
            
            task_manager.update_state(
                state=task_state["state"],  # 保持原有状态
                output_result=output_result
            )
            
            return True
        except Exception as e:
            self.logger.error(f"保存回顾结果失败: task_id={task_id}, error={str(e)}")
            return False

# 全局任务回顾管理器实例
task_review_manager = TaskReviewManager()

# 测试代码
async def test_task_review():
    """测试任务回顾管理器"""
    logger.info("开始测试任务回顾管理器")
    
    # 模拟任务ID
    test_task_id = "test_task_review_001"
    
    # 测试回顾功能
    result = await task_review_manager.review_task_results(test_task_id, model_backend="vllm")
    logger.info(f"任务回顾测试结果: {result}")
    
    logger.info("任务回顾管理器测试完成")

if __name__ == "__main__":
    import asyncio
    asyncio.run(test_task_review())
